import atexit
import threading
from typing import Union, Dict, List

import mysql
from mysql.connector.abstracts import MySQLConnectionAbstract
from mysql.connector.pooling import PooledMySQLConnection
from mysql.connector.types import RowItemType, RowType

import utils
from utils import load_json

conns: dict[str: Union[PooledMySQLConnection, MySQLConnectionAbstract]] = {}

db_conn_lock: threading.Lock = threading.Lock()


def is_alive(conn: Union[PooledMySQLConnection, MySQLConnectionAbstract]):
    return conn and conn.is_connected()


def get_mysql_conn(config_json: str = "../config/mysql.json") -> Union[PooledMySQLConnection, MySQLConnectionAbstract]:
    config = load_json(config_json)
    k = utils.md5(config)
    try:
        db_conn_lock.acquire()
        conn = conns.get(k)
        if not is_alive(conn):
            conn = mysql.connector.connect(**config)
            conns[k] = conn
        return conn
    finally:
        db_conn_lock.release()


def release_mysql_conn(k: str, conn: Union[PooledMySQLConnection, MySQLConnectionAbstract]):
    if is_alive(conn):
        conn.close()
        print(f"连接释放,id:{k}")


def mysql_query(conn, sql: str) -> List[Union[RowType, Dict[str, RowItemType]]]:
    cursor = conn.cursor()
    cursor.execute(sql)
    results = cursor.fetchall()
    cursor.close()
    return results


def release_conns():
    print("release conns start.")
    for k, c in conns.items():
        release_mysql_conn(k, c)
    print("release conns finished.")


def register_hook():
    atexit.register(release_conns)
    print("register_hook")


if __name__ == '__main__':
    conn = get_mysql_conn()
    res = mysql_query(conn, "select version()")
    print(res)
